use std::collections::{HashMap, HashSet};
use std::sync::{Arc, Mutex};
use serde::{Deserialize, Serialize};
use tokio::sync::broadcast;
use futures_util::{SinkExt, StreamExt};
use tokio::net::{TcpListener, TcpStream};
use tokio_tungstenite::{accept_async, tungstenite::Message, WebSocketStream};
use uuid::Uuid;

use crate::utils::message_queue::MessageQueue;

// 客户端消息
#[derive(Debug, Clone, Deserialize)]
pub struct ClientMessage {
    pub r#type: String,
    pub topic: Option<String>,
    pub data: Option<serde_json::Value>,
}

// 服务器消息
#[derive(Debug, Clone, Serialize)]
pub struct ServerMessage {
    pub topic: String,
    pub data: serde_json::Value,
    pub timestamp: u64,
}

// 心跳消息
#[derive(Debug, Clone, Serialize)]
pub struct HeartbeatMessage {
    pub r#type: String,
    pub timestamp: u64,
}

// WebSocket连接
struct Connection {
    id: String,
    subscriptions: HashSet<String>,
}

// WebSocket服务
pub struct WebSocketService {
    clients: Arc<Mutex<HashMap<String, WebSocketStream<TcpStream>>>>,
    connections: Arc<Mutex<HashMap<String, Connection>>>,
    message_sender: broadcast::Sender<ServerMessage>,
    message_queue: Arc<Mutex<MessageQueue>>,
}

impl WebSocketService {
    pub fn new(message_queue: Arc<Mutex<MessageQueue>>) -> Self {
        let (sender, _) = broadcast::channel(100);
        
        Self {
            clients: Arc::new(Mutex::new(HashMap::new())),
            connections: Arc::new(Mutex::new(HashMap::new())),
            message_sender: sender,
            message_queue,
        }
    }

    // 启动WebSocket服务器
    pub async fn start(&self, addr: &str) -> Result<(), Box<dyn std::error::Error>> {
        let listener = TcpListener::bind(addr).await?;
        println!("WebSocket服务器已启动，监听地址: {}", addr);

        let clients = self.clients.clone();
        let connections = self.connections.clone();
        let message_sender = self.message_sender.clone();
        let message_queue = self.message_queue.clone();

        // 启动消息分发任务
        self.start_message_dispatcher();

        // 接受连接
        while let Ok((stream, addr)) = listener.accept().await {
            println!("新的连接: {}", addr);

            let ws_stream = accept_async(stream).await?;
            let client_id = Uuid::new_v4().to_string();
            
            // 创建连接记录
            {
                let mut connections = connections.lock().unwrap();
                connections.insert(client_id.clone(), Connection {
                    id: client_id.clone(),
                    subscriptions: HashSet::new(),
                });
            }

            // 存储客户端连接
            {
                let mut clients = clients.lock().unwrap();
                clients.insert(client_id.clone(), ws_stream);
            }

            // 处理客户端连接
            let client_connections = connections.clone();
            let client_clients = clients.clone();
            let client_sender = message_sender.clone();
            let client_queue = message_queue.clone();

            tokio::spawn(async move {
                Self::handle_connection(
                    client_id,
                    client_connections,
                    client_clients,
                    client_sender,
                    client_queue
                ).await;
            });
        }

        Ok(())
    }

    // 处理客户端连接
    async fn handle_connection(
        client_id: String,
        connections: Arc<Mutex<HashMap<String, Connection>>>,
        clients: Arc<Mutex<HashMap<String, WebSocketStream<TcpStream>>>>,
        message_sender: broadcast::Sender<ServerMessage>,
        message_queue: Arc<Mutex<MessageQueue>>,
    ) {
        // 获取客户端连接
        let mut client_stream = {
            let mut clients = clients.lock().unwrap();
            clients.remove(&client_id).unwrap()
        };

        // 创建消息接收器
        let mut message_receiver = message_sender.subscribe();

        // 处理消息循环
        loop {
            tokio::select! {
                // 处理来自客户端的消息
                Some(msg) = client_stream.next() => {
                    match msg {
                        Ok(Message::Text(text)) => {
                            if let Err(e) = Self::handle_client_message(
                                &client_id,
                                &text,
                                &connections,
                                &message_queue
                            ).await {
                                println!("处理客户端消息错误: {}", e);
                            }
                        },
                        Ok(Message::Close(_)) => {
                            println!("客户端断开连接: {}", client_id);
                            break;
                        },
                        Err(e) => {
                            println!("WebSocket错误: {}", e);
                            break;
                        },
                        _ => {}
                    }
                },
                
                // 处理广播消息
                Ok(server_msg) = message_receiver.recv() => {
                    // 检查客户端是否订阅了该主题
                    let is_subscribed = {
                        let connections = connections.lock().unwrap();
                        if let Some(connection) = connections.get(&client_id) {
                            connection.subscriptions.contains(&server_msg.topic)
                        } else {
                            false
                        }
                    };

                    // 如果订阅了，发送消息
                    if is_subscribed {
                        let json = serde_json::to_string(&server_msg).unwrap();
                        if let Err(e) = client_stream.send(Message::Text(json)).await {
                            println!("发送消息错误: {}", e);
                            break;
                        }
                    }
                }
            }
        }

        // 清理连接
        {
            let mut connections = connections.lock().unwrap();
            connections.remove(&client_id);
        }
    }

    // 处理客户端消息
    async fn handle_client_message(
        client_id: &str,
        message: &str,
        connections: &Arc<Mutex<HashMap<String, Connection>>>,
        message_queue: &Arc<Mutex<MessageQueue>>,
    ) -> Result<(), Box<dyn std::error::Error>> {
        // 解析消息
        let client_msg: ClientMessage = serde_json::from_str(message)?;

        match client_msg.r#type.as_str() {
            // 订阅主题
            "subscribe" => {
                if let Some(topic) = client_msg.topic {
                    let mut connections = connections.lock().unwrap();
                    if let Some(connection) = connections.get_mut(client_id) {
                        connection.subscriptions.insert(topic.clone());
                        println!("客户端 {} 订阅主题: {}", client_id, topic);
                    }
                }
            },
            
            // 取消订阅
            "unsubscribe" => {
                if let Some(topic) = client_msg.topic {
                    let mut connections = connections.lock().unwrap();
                    if let Some(connection) = connections.get_mut(client_id) {
                        connection.subscriptions.remove(&topic);
                        println!("客户端 {} 取消订阅主题: {}", client_id, topic);
                    }
                }
            },
            
            // 心跳
            "heartbeat" => {
                // 处理心跳，可以记录客户端活动时间
                println!("收到客户端 {} 的心跳", client_id);
            },
            
            // 客户端发送数据
            "data" => {
                if let (Some(topic), Some(data)) = (client_msg.topic, client_msg.data) {
                    // 将数据推送到消息队列
                    let mut queue = message_queue.lock().unwrap();
                    queue.push_message(topic, data.to_string())?;
                }
            },
            
            _ => {
                println!("未知消息类型: {}", client_msg.r#type);
            }
        }

        Ok(())
    }

    // 启动消息分发任务
    fn start_message_dispatcher(&self) {
        let message_queue = self.message_queue.clone();
        let sender = self.message_sender.clone();

        tokio::spawn(async move {
            loop {
                // 每100ms检查一次消息队列
                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
                
                // 获取所有主题的消息
                let topics = {
                    let queue = message_queue.lock().unwrap();
                    queue.get_topics()
                };

                for topic in topics {
                    // 获取主题的消息
                    let messages = {
                        let mut queue = message_queue.lock().unwrap();
                        queue.get_messages(&topic, 10) // 每次最多获取10条消息
                    };

                    // 发送消息
                    for msg in messages {
                        let timestamp = std::time::SystemTime::now()
                            .duration_since(std::time::UNIX_EPOCH)
                            .unwrap()
                            .as_secs();

                        let server_msg = ServerMessage {
                            topic: topic.clone(),
                            data: serde_json::from_str(&msg).unwrap_or(serde_json::json!({})),
                            timestamp,
                        };

                        // 广播消息
                        if let Err(e) = sender.send(server_msg) {
                            println!("广播消息错误: {}", e);
                        }
                    }
                }
            }
        });
    }
}